import os
import numpy as np
from glob import glob
import json
import cv2
import random
from copy import deepcopy
import math
from shapely.geometry import Polygon

try:
    from .cvio import cvio
except Exception as e:
    from cvio import cvio

class MosaicAffineData(object):
    def __init__(self, data_root, mosaic_size=800, mosaic_prob=0.5, affine_prob=0.5, degree=30, scale=0.1):
        # self.image_list = glob(f'{data_root}{os.path.sep}*.jp*')
        # self.label_list = [f'{im.split(".jp")[0]}.json' for im in self.image_list]
        self.img_ann_list = cvio.load_img_ann_list(data_root)
        self.image_list, self.label_list = [], []
        for img, ann in self.img_ann_list:
            self.image_list.append(img)
            self.label_list.append(ann)
        self.mosaic_size = mosaic_size
        self.mosaic_border = [-mosaic_size // 2, -mosaic_size // 2]
        self.mosaic_prob = mosaic_prob
        self.affine_prob = affine_prob
        self.degree = degree
        self.scale = scale

    def shape2bbox(self, shape):
        bbox = np.array(shape['points'])
        x1 = bbox[:, 0].min()
        x2 = bbox[:, 0].max()
        y1 = bbox[:, 1].min()
        y2 = bbox[:, 1].max()
        bbox = [x1, y1, x2, y2]
        return bbox

    def load_at(self, index):
        img = cv2.imread(self.image_list[index])
        with open(self.label_list[index]) as fp:
            label = json.load(fp)
        shapes = label['shapes']
        for shape in shapes:
            shape_type = shape['shape_type']
            if shape_type == 'rectangle':
                x1, y1, x2, y2 = self.shape2bbox(shape)
                shape['points'] = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]
                shape['shape_type'] = 'polygon'
        return img, label

    def create_mosaic(self, index, rm_thr=0.3):
        s = self.mosaic_size
        yc, xc = [int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border]  # mosaic center x, y
        indices = [index] + [random.randint(0, len(self.image_list) - 1) for _ in range(3)]  # 3 additional image indices
        img, label = self.load_at(index)
        label4 = deepcopy(label)
        label4['shapes'] = []
        postfix = ''
        for i, index in enumerate(indices):
            img, label = self.load_at(index)
            h, w = label['imageHeight'], label['imageWidth']
            # place img in img4
            if i == 0:  # top left
                img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8)  # base image with 4 tiles
                x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc  # xmin, ymin, xmax, ymax (large image)
                x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h  # xmin, ymin, xmax, ymax (small image)
                dwh = (xc - w, yc - h)
            elif i == 1:  # top right
                x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc
                x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
                dwh = (xc, yc - h)
            elif i == 2:  # bottom left
                x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)
                x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, max(xc, w), min(y2a - y1a, h)
                dwh = (xc - w, yc)
            elif i == 3:  # bottom right
                x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)
                x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)
                dwh = (xc, yc)

            img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b]  # img4[ymin:ymax, xmin:xmax]
            label4 = self.translate_points(label, label4, dwh, rm_thr)
            postfix += f'_{index}'

        # 更改json中图片宽高和图片保存路径
        postfix = '_mosaic'
        imageName = os.path.basename(self.image_list[index]).split('.jp')[0]
        imageName += postfix
        label4['imagePath'] = f'{imageName}.jpg'
        label4['imageWidth'] = s * 2
        label4['imageHeight'] = s * 2
        return img4, label4

    def translate_points(self, label, label4, dwh, rm_thr=0.3):
        label = deepcopy(label)
        ms = self.mosaic_size * 2
        shapes = label['shapes']
        for shape in shapes:
            points = shape['points']
            ori_area = Polygon(np.array(points)).area
            for point in points:
                point[0] += dwh[0]
                point[1] += dwh[1]
                for i, p in enumerate(point):
                    if p < 0 :
                        point[i] = 0
                    elif p > ms:
                        point[i] = ms
            new_area = Polygon(np.array(points)).area
            if new_area > ori_area * rm_thr:
                label4['shapes'].append(deepcopy(shape))
        return label4


    def random_affine(self, img, label, angle, scale, translate=.1, shear=10, border=(0, 0), rm_thr=0.1):

        height = img.shape[0] + border[0] * 2  # shape(h,w,c)
        width = img.shape[1] + border[1] * 2

        # Rotation and Scale
        R = np.eye(3)
        R[:2] = cv2.getRotationMatrix2D(angle=angle, center=(img.shape[1] / 2, img.shape[0] / 2), scale=scale)

        # Translation
        T = np.eye(3)
        T[0, 2] = random.uniform(-translate, translate) * img.shape[1] + border[1]  # x translation (pixels)
        T[1, 2] = random.uniform(-translate, translate) * img.shape[0] + border[0]  # y translation (pixels)

        # Shear
        S = np.eye(3)
        S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # x shear (deg)
        S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # y shear (deg)

        # Combined rotation matrix
        M = S @ T @ R  # 顺序很重要!!
        if (border[0] != 0) or (border[1] != 0) or (M != np.eye(3)).any():  # image changed
            img = cv2.warpAffine(img, M[:2], dsize=(width, height), flags=cv2.INTER_LINEAR, borderValue=(114, 114, 114))

        h, w = img.shape[:2]
        shapes = label['shapes']
        new_shapes = []
        for shape in shapes:
            points = shape['points']
            ori_area = Polygon(np.array(points)).area
            for i, point in enumerate(points):
                point = np.dot(M, np.array(point + [1,]))[:-1]
                if point[0] < 0:
                    point[0] = 0
                elif point[0] > w:
                    point[0] = w
                if point[1] < 0:
                    point[1] = 0
                elif point[1] > h:
                    point[1] = h
                shape['points'][i] = point.tolist()[:2]
            new_area = Polygon(np.array(points)).area
            if new_area > ori_area * rm_thr:
                new_shapes.append(shape)
            # new_shapes.append(shape)
        label["shapes"] = new_shapes
        label['imagePath'] = label['imagePath'].split('.jp')[0] + f'_affine.jpg'
        return img, label

    def remove_small_polygons(self, ori_label, new_label, rm_thr=100, area_thr=0.3):
        shapes = []
        for i, (oshape, nshape) in enumerate(zip(ori_label['shapes'], new_label['shapes'])):
            opoints = np.array(oshape['points'])
            npoints = np.array(nshape['points'])
            if len(opoints) <= 2 or len(npoints) <= 2:
                continue
            oarea = Polygon(opoints).area
            narea = Polygon(npoints).area
            if narea < rm_thr or narea < oarea*area_thr:
                continue
            shapes.append((nshape))
        if len(shapes) == 0:
            return None
        new_label['shapes'] = shapes
        return new_label

    def createAugmentData(self, save_root, image_list=None, rm_thr=0.3):
        if image_list is not None:
            self.image_list = [i for i in image_list if i in msafData.image_list]
            self.label_list = [f'{im.split(".jp")[0]}.json' for im in self.image_list]
        if os.path.exists(save_root):
            import shutil
            shutil.rmtree(save_root)
        os.makedirs(save_root)
        tbar = range(len(self.image_list))
        
        for i in tbar:
            img, label = self[i]
            print(self.image_list[i])
            if img is None:
                print(f'Image skipped (not exist): {self.image_list[i]}')
                continue
            is_mosaic = random.random() < self.mosaic_prob
            is_affine = random.random() < self.affine_prob
            is_aug = is_mosaic or is_affine
            if not is_aug:
                continue
            if is_mosaic:
                img, label = self.create_mosaic(i, rm_thr)
            if is_affine:
                angle = random.uniform(-self.degree, self.degree)
                scale = random.uniform(1 - self.scale, 1 + self.scale)
                img, label = self.random_affine(img, label, angle, scale, translate=.01, shear=10, border=(0, 0))

            if label is None:
                print(f'Dirty annotation skipped(some shapes have less than 3 points): {self.image_list[i]}')
                continue
            label['imageData'] = None
            self.save_at(img, label, save_root)

    def __len__(self):
        return len(self.image_list)

    def __getitem__(self, item):
        return self.load_at(item)

    def save_at(self, img, label, save_root, imageName=None):
        h, w, c = img.shape
        label['imageHeight'] = h
        label['imageWidth'] = w
        if imageName is None:
            imageName = label['imagePath'].split('.jpg')[0]
        if '.jp' in imageName:
            imageName = imageName.split('.jp')[0]

        cv2.imwrite(f"{save_root}{os.path.sep}{imageName}.jpg", img)
        with open(f"{save_root}{os.path.sep}{imageName}.json", 'w') as fp:
            json.dump(label, fp)

def check_data_from_coco(ann_path, skuNumThr=5000):
    # 1.查看数据集图片数量
    # 2.查看sku数量
    # 3.返回数量较少的sku所对应的图片列表，便于后续进行增强
    with open(ann_path) as fp:
        annDict = json.load(fp)
    annotations = annDict['annotations']
    images = annDict['images']
    categories = annDict['categories']
    sku_stats = {n['name']: [] for n in categories}
    for ann in annotations:
        image_id = ann['image_id']
        file_name = images[image_id]['file_name']
        category_id = ann['category_id']
        category = categories[category_id]['name']
        sku_stats[category].append(file_name)

    info = ''
    for k, v in sku_stats.items():
        num = len(v)
        info += f'{k}   {num}\n'
    print(info)
    sku_stats = {k: set(v) for k, v in sku_stats.items() if len(v) < skuNumThr}
    return sku_stats


def create_sku_list(root, sku_stats):
    image_list = []
    for k, v in sku_stats.items():
        for i in v:
            image_list.append(f'{root}/{i.replace("JPEGImages/", "")}')
    image_list = list(set(image_list))
    return image_list

def createMosaicAffineDataMask(data_root, save_root, mosaic_size=800, 
            mosaic_prob=1, affine_prob=0, degree=15, scale=1, rm_thr=0.5):
    msafData = MosaicAffineData(data_root, mosaic_size=mosaic_size, mosaic_prob=mosaic_prob, 
            affine_prob=affine_prob, degree=degree, scale=scale)
    msafData.createAugmentData(save_root=save_root, rm_thr=rm_thr)        



if __name__ == '__main__':
    idx = 1
    data_root = r'F:\Dataset\jiujiang\modeling_data\labelme\temp\base'#\%d' % idx
    save_root = r'F:\Dataset\jiujiang\modeling_data\labelme\temp\mosaic\%d' % idx
    msafData = MosaicAffineData(data_root, mosaic_size=800, mosaic_prob=1, affine_prob=0.0, degree=15, scale=1)
    msafData.createAugmentData(save_root=save_root, rm_thr=0.5)

    # 根据训练集json文件统计图片数量较少的sku, 返回其对应的图片列表
    # sku_stats = check_data_from_coco(ann_path='/mnt/g/超级模型数据/daiding5.0/103_rh/rh_hr/annotations/train.json',
    #                                  skuNumThr=300)
    # image_list = create_sku_list(data_root, sku_stats)
    # 替换掉msafData对象里边的图片和标签列表
    # 生成mosaic增广图片
    # msafData.createAugmentData(save_root=save_root, rm_thr=0.25, image_list=image_list)


